package com.llmops.core.workflow;

import com.alibaba.fastjson.JSONObject;
import com.llmops.core.Edge;
import com.llmops.core.Node;
import com.llmops.core.WorkflowGraph;
import com.llmops.core.node.*;

import java.util.ArrayList;
import java.util.List;

public class WorkflowParser {
    
    public Node parseNode(NodeConfig nf) {
        switch (nf.getType()) {
            case "userInputNode":
                return new StartNode(nf.getId(),nf.getTitle(),nf.getConfig());
            case "llmNode":
                return new LLMNode(nf.getId(),nf.getTitle(),nf.getConfig());
            case "classificationNode":
                return new ClassificationNode(nf.getId(),nf.getTitle(),nf.getConfig());
            case "ragNode":
                return new RagNode(nf.getId(),nf.getTitle(),nf.getConfig());
            case "outputNode":
                return new EndNode(nf.getId(),nf.getTitle(),nf.getConfig());
            default:
                throw new IllegalArgumentException("未知节点类型: " + nf.getType());
        }
    }

    public WorkflowGraph parserFlow(String json){
        WorkflowGraph workflowGraph = new WorkflowGraph("1");
        Workflow workflow = JSONObject.parseObject(json, Workflow.class);
        List<NodeConfig> nfs = workflow.getNodes();
        for (NodeConfig nf : nfs) {
            workflowGraph.addNode(parseNode(nf));
        }

        List<Connection> cfs = workflow.getConnections();
        for (Connection cf : cfs) {
            workflowGraph.addEdge(cf.getSourceId(),cf.getTargetId(),
                    cf.getConditionType()==null?null:Edge.ConditionType.valueOf(cf.getConditionType()),cf.getConditionValue());
        }
        return workflowGraph;
    }

    public String parserFlow(WorkflowGraph workflowGraph){
        Workflow workflow = new Workflow();
        workflow.setId(workflowGraph.getId());
        List<NodeConfig> nodes = new ArrayList<>();
        List<Connection> connections = new ArrayList<>();
        for (Node value : workflowGraph.getNodes().values()) {
            NodeConfig nodeConfig = new NodeConfig(value.getId(),value.getType(),value.getConfig());
            nodes.add(nodeConfig);
        }
        for (Edge edge : workflowGraph.getEdges()) {
            Connection connection = new Connection(edge.getSourceId(),edge.getTargetId());
            connections.add(connection);
        }
        workflow.setNodes(nodes);
        workflow.setConnections(connections);
        return JSONObject.toJSONString(workflow);
    }

}